package org.homeworkteam.webSocket;

import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.concurrent.ConcurrentHashMap;

@Component
public class SimpleWebSocketHandler extends TextWebSocketHandler {

    // 使用线程安全的 Map 来存储用户和 WebSocket 会话
    private static final ConcurrentHashMap<String, WebSocketSession> userSessions = new ConcurrentHashMap<>();

    /**
     * 建立连接后保存用户的 WebSocket 会话
     *
     * @param session WebSocket 会话
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        try {
            // 从连接 URI 中解析 userId
            String userId = session.getUri().getQuery().split("=")[1];
            userSessions.put(userId, session);
            System.out.println("User connected: " + userId);
        } catch (Exception e) {
            System.out.println("Failed to extract userId from WebSocket connection: " + e.getMessage());
        }
    }

    /**
     * 连接关闭时移除会话
     *
     * @param session WebSocket 会话
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) {
        userSessions.values().remove(session);
        System.out.println("Connection closed: " + session.getId());
    }

    /**
     * 发送消息给指定用户
     *
     * @param userId  接收者 ID
     * @param message 消息内容
     * @return 是否发送成功
     */
    public boolean sendMessageToUser(String userId, String message) {
        WebSocketSession session = userSessions.get(userId); // 获取目标用户的 WebSocket 会话
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(message)); // 推送消息
                System.out.println("Message sent to user " + userId + ": " + message);
                return true;
            } catch (Exception e) {
                System.out.println("Failed to send message to user " + userId + ": " + e.getMessage());
                return false;
            }
        } else {
            System.out.println("User " + userId + " is not connected.");
            return false;
        }
    }
}
